package com.embeddings.embed;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import okhttp3.*;

import java.io.IOException;
import java.util.*;

public class EmbedClient {

    public static double[] getEmbedding(String uri, String apiKey, String inputText) throws IOException {
        OkHttpClient client = new OkHttpClient();

        // 创建请求体
        JSONObject requestBody = new JSONObject();
        requestBody.put("input", Collections.singletonList(inputText));

        // 创建请求
        MediaType mediaType = MediaType.parse("application/json; charset=utf-8");
        RequestBody body = RequestBody.Companion.create(requestBody.toJSONString(), mediaType);
        Request request = new Request.Builder()
                .url(uri)
                .addHeader("Authorization", "Bearer " + apiKey)
                .addHeader("Content-Type", "application/json")
                .post(body)
                .build();

        // 发送请求
        Response response = client.newCall(request).execute();
        if (!response.isSuccessful()) {
            throw new IOException("Unexpected code " + response);
        }

        // 解析JSON响应
        String responseBody = response.body().string();
        EmbeddingResponse embeddingResponse = JSON.parseObject(responseBody, EmbeddingResponse.class);

        // 返回嵌入向量
        return embeddingResponse.getData().get(0).getEmbedding();
    }

    static class EmbeddingResponse {
        private List<Data> data;

        public List<Data> getData() {
            return data;
        }

        public void setData(List<Data> data) {
            this.data = data;
        }
    }

    static class Data {
        private double[] embedding;
        private int index;
        private String object;

        public double[] getEmbedding() {
            return embedding;
        }

        public void setEmbedding(double[] embedding) {
            this.embedding = embedding;
        }

        public int getIndex() {
            return index;
        }

        public void setIndex(int index) {
            this.index = index;
        }

        public String getObject() {
            return object;
        }

        public void setObject(String object) {
            this.object = object;
        }
    }
}